'''
正弦余弦算法

'''
import numpy as np
import random
import math
import matplotlib.pyplot as plt
import time
from chart3d import Chart3D
import obj_funs


class SCA:
    def __init__(self, n_dim=2, pop_size=10, max_iter=100, lb=[0], ub=[50], obj_func=None, a=2, ax=None):
        self.pop = pop_size
        self.n_dim = n_dim
        self.a = a  # 感知概率
        self.func = obj_func
        self.max_iter = max_iter  # max iter
        self.ax = ax  # chart对象，用于3d绘图描点
        self.lb, self.ub = np.array(lb) * np.ones(self.n_dim), np.array(ub) * np.ones(self.n_dim)
        assert self.n_dim == len(self.lb) == len(self.ub), 'dim == len(lb) == len(ub) is not True'
        assert np.all(self.ub > self.lb), 'upper-bound must be greater than lower-bound'

        # self.X = np.random.uniform(low=self.lb, high=self.ub, size=(self.pop, self.n_dim))
        self.X = np.zeros([pop_size, n_dim])
        for i in range(pop_size):
            for j in range(n_dim):
                self.X[i, j] = np.random.random() * (ub[j] - lb[j]) + lb[j]
        self.Y = [self.func(self.X[i]) for i in range(len(self.X))]
        if self.ax:
            # 初始化标点
            self.ax.chat3dplot(self.X, self.func, 'ax1', 'o', 'black')

        self.pbest_x = self.X.copy()  # 历史上每一个粒子的个人最佳位置
        self.pbest_y = [np.inf for i in range(self.pop)]  # 历史上每一个粒子的最佳适应值
        self.gbest_x = self.pbest_x.mean(axis=0).reshape(1, -1)  # 所有粒子的全局最佳位置
        self.gbest_y = np.inf  # 所有粒子的全局最佳值
        self.gbest_y_hist = []  # 每次迭代最佳值数组
        self.update_pbest()
        self.update_gbest()
        self.gbest_y_hist.append(self.gbest_y)

    def update_pbest(self):
        '''
        personal best
        :return:
        '''
        for i in range(len(self.Y)):
            if self.pbest_y[i] > self.Y[i]:
                self.pbest_x[i] = self.X[i]
                self.pbest_y[i] = self.Y[i]

    def update_gbest(self):
        '''
        global best
        :return:
        '''
        idx_min = self.pbest_y.index(min(self.pbest_y))
        if self.gbest_y > self.pbest_y[idx_min]:
            self.gbest_x = self.X[idx_min, :].copy()
            self.gbest_y = self.pbest_y[idx_min]

    def update(self, i):
        r1 = self.a - i * ((self.a) / self.max_iter)
        for j in range(self.pop):
            for k in range(self.n_dim):
                r2 = 2 * math.pi * random.uniform(0.0, 1.0)
                r3 = 2 * random.uniform(0.0, 1.0)
                r4 = random.uniform(0.0, 1.0)
                if r4 < 0.5:
                    try:
                        self.X[j][k] = self.X[j][k] + (r1 * math.sin(r2) * abs(r3 * self.gbest_x[k] - self.X[j][k]))
                    except:
                        self.X[j][k] = self.X[j][k] + (r1 * math.sin(r2) * abs(r3 * self.gbest_x[0][k] - self.X[j][k]))
                else:
                    try:
                        self.X[j][k] = self.X[j][k] + (r1 * math.cos(r2) * abs(r3 * self.gbest_x[k] - self.X[j][k]))
                    except:
                        self.X[j][k] = self.X[j][k] + (r1 * math.cos(r2) * abs(r3 * self.gbest_x[0][k] - self.X[j][k]))
        self.X = np.clip(self.X, self.lb, self.ub)
        self.Y = [self.func(self.X[i]) for i in range(len(self.X))]  # Function for fitness evaluation of new solutions

    def sca(self):
        avg_fit = []
        avg_ts = []  # 平均方差
        time_start = time.time()  # 记录迭代寻优开始时间
        for i in range(self.max_iter):
            self.update(i)
            avg_fit.append(np.mean(self.Y))
            avg_ts.append(np.var(self.Y))
            self.update_pbest()
            self.update_gbest()
            self.gbest_y_hist.append(self.gbest_y)

            # 展示描点图
            if self.ax and i == self.ax.a1:
                # 最终结果化标点
                self.ax.chat3dplot(self.X, self.func, 'ax2', 'o', 'black')

            elif self.ax and i == self.ax.a2:
                self.ax.chat3dplot(self.X, self.func, 'ax3', 'o', 'black')

        time_end = time.time()  # 记录迭代结束时间
        print(f'SCA共花费 {time_end - time_start} 秒')
        print('SCA最优适应度', self.gbest_y)
        print('SCA最优解', self.gbest_x)

        if (self.ax):
            # 最终结果化标点
            self.ax.chat3dplot(self.X, self.func, 'ax4', 'o', 'black')
            plt.show()

        self.best_x, self.best_y = self.gbest_x, self.gbest_y

        return self.gbest_y_hist, avg_fit, avg_ts
